package com.lsa.letcode.balanced_binary_tree;

import java.util.ArrayList;
import java.util.EmptyStackException;

import com.lsa.letcode.common.TreeNode;

public class Solution {
	private static Stack<StackElement> stack = new Stack<StackElement>();
    public boolean isBalanced(TreeNode r) {
        return heightIterative(r) + 1 > -1;
    }
    
    static int height(TreeNode r) {
    	if (r == null) {return -1;} // line 0
    	
    	int leftHeight = height(r.left); // line 1
    	int rightHeight = height(r.right); // line 2
    	if (leftHeight == -2 || rightHeight == -2  || Math.abs(leftHeight - rightHeight) > 1) return -2; // line 3
    	return Math.max(leftHeight, rightHeight) + 1;
    }
    
    private static class StackElement {
    	TreeNode r;
    	
    	int lineNo = 0;
    	int leftHeight;
    	int rightHeight;
    	
		public StackElement(TreeNode r) {
			this.r = r;
		}

		@Override
		public String toString() {
			return "StackElement [r=" + r + ", lineNo=" + lineNo + "]";
		}
    }
    
    static int heightIterative(TreeNode r) {
    	stack.clear();
    	stack.add(new StackElement(r));
    	
    	int returnValue = 0;
    	while(!stack.isEmpty()) {
    		StackElement top = stack.peek();
    		int lineNo = top.lineNo++;
    		if (lineNo == 0) {
    			if (top.r == null) {
    				returnValue = -1;
    				stack.pop();
    			}
    		}
    		else if (lineNo == 1) {
    			stack.add(new StackElement(top.r.left));
    		}
    		else if (lineNo == 2) {
    			top.leftHeight = returnValue;
    			stack.add(new StackElement(top.r.right));
    		}
    		else {
    			top.rightHeight = returnValue;
    			if (top.leftHeight == -2 || top.rightHeight  == -2  || Math.abs(top.leftHeight - top.rightHeight) > 1) {
    				returnValue = -2;
    			}
    			else {
    				returnValue = Math.max(top.leftHeight, top.rightHeight) + 1;
    			}
    			stack.pop();
    		}
    	}
    	return returnValue;
    }
    
    @SuppressWarnings("serial")
	public static class Stack<E> extends ArrayList<E> {
        /**
         * Creates an empty Stack.
         */
        public Stack() {
        	super(1024*1024);
        }

        /**
         * Removes the object at the top of this stack and returns that
         * object as the value of this function.
         *
         * @return  The object at the top of this stack (the last item
         *          of the <tt>Vector</tt> object).
         * @throws  EmptyStackException  if this stack is empty.
         */
        public E pop() {
            E       obj;
            int     len = size();
            obj = peek();
            remove(len - 1);

            return obj;
        }

        /**
         * Looks at the object at the top of this stack without removing it
         * from the stack.
         *
         * @return  the object at the top of this stack (the last item
         *          of the <tt>Vector</tt> object).
         * @throws  EmptyStackException  if this stack is empty.
         */
        public E peek() {
            return get(size() - 1);
        }

        /**
         * Tests if this stack is empty.
         *
         * @return  <code>true</code> if and only if this stack contains
         *          no items; <code>false</code> otherwise.
         */
        public boolean empty() {
            return size() == 0;
        }
    }
}